#include <stdint.h>
#include <immintrin.h>
#include "utils.h"

int l1(int niters, int8_t *A, int8_t *B, int32_t *C, size_t M, size_t N, size_t K) {
  for(int i = 0; i < niters; ++i) {
    for(int mi=0; mi<(M/TILE_M); mi+=2) {
      for(int ni=0; ni<(N/TILE_N); ni+=2) {
        _tile_zero(TMM4);
        _tile_zero(TMM5);
        _tile_zero(TMM6);
        _tile_zero(TMM7);
        for(int ki=0; ki<(K/TILE_K); ++ki) {
          // Load A0
          _tile_loadd(TMM0, A+mi*TILE_M*K+ki*TILE_K, K);
          // _tile_stream_loadd(TMM0, A+mi*TILE_M*K+ki*TILE_K, K);
          // Load B0
          _tile_loadd(TMM2, B+ki*(TILE_K/4)*(N*4)+ni*TILE_N*4, N*4);
          // Load B0 (bypass L1)
          // _tile_stream_loadd(TMM2, B+ki*(TILE_K/4)*(N*4)+ni*TILE_N*4, N*4);
          // C00 += A0 * B0
          _tile_dpbuud(TMM4, TMM0, TMM2);
          // Load A1
          _tile_loadd(TMM1, A+(mi+1)*TILE_M*K+ki*TILE_K, K);
          // _tile_stream_loadd(TMM1, A+(mi+1)*TILE_M*K+ki*TILE_K, K);
          // C01 += A0 * B1
          _tile_dpbuud(TMM6, TMM1, TMM2);
          // Load B1
          _tile_loadd(TMM3, B+ki*(TILE_K/4)*(N*4)+(ni+1)*TILE_N*4, N*4);
          // Load B1 (bypass L1)
          // _tile_stream_loadd(TMM3, B+ki*(TILE_K/4)*(N*4)+(ni+1)*TILE_N*4, N*4);
          // C10 += A1 * B0
          _tile_dpbuud(TMM5, TMM0, TMM3);
          // C11 += A1 * B1
          _tile_dpbuud(TMM7, TMM1, TMM3);
        }
        // Store C00
        _tile_stored(TMM4, (int8_t *)C+mi*TILE_M*N*sizeof(int32_t)+ni*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        // Store C01
        _tile_stored(TMM5, (int8_t *)C+mi*TILE_M*N*sizeof(int32_t)+(ni+1)*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        // Store C10
        _tile_stored(TMM6, (int8_t *)C+(mi+1)*TILE_M*N*sizeof(int32_t)+ni*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        // Store C11
        _tile_stored(TMM7, (int8_t *)C+(mi+1)*TILE_M*N*sizeof(int32_t)+(ni+1)*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
      }
    }
  }
  return 0;
}


// load bw utilization: 50% (limited by store)
// store bw utilization: 100%
int l1_bw_load_store(int niters, int8_t *A, int8_t *B, int32_t *C, size_t M, size_t N, size_t K) {
  for(int i = 0; i < niters; ++i) {
    for(int mi=0; mi<(M/TILE_M); mi+=2) {
      for(int ni=0; ni<(N/TILE_N); ni+=2) {
        for(int ki=0; ki<(K/TILE_K); ++ki) {
          _tile_loadd(TMM0, A+mi*TILE_M*K+ki*TILE_K, K);
          _tile_loadd(TMM2, B+ki*(TILE_K/4)*(N*4)+ni*TILE_N*4, N*4);
          _tile_loadd(TMM1, A+(mi+1)*TILE_M*K+ki*TILE_K, K);
          _tile_loadd(TMM3, B+ki*(TILE_K/4)*(N*4)+(ni+1)*TILE_N*4, N*4);
        }
        _tile_zero(TMM0);
        _tile_zero(TMM1);
        _tile_zero(TMM2);
        _tile_zero(TMM3);

        _tile_stored(TMM0, (int8_t *)C+mi*TILE_M*N*sizeof(int32_t)+ni*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        _tile_stored(TMM1, (int8_t *)C+mi*TILE_M*N*sizeof(int32_t)+(ni+1)*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        _tile_stored(TMM2, (int8_t *)C+(mi+1)*TILE_M*N*sizeof(int32_t)+ni*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        _tile_stored(TMM3, (int8_t *)C+(mi+1)*TILE_M*N*sizeof(int32_t)+(ni+1)*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
      }
    }
  }
  return 0;
}

// load bw utilization: 100%
int l1_bw_load(int niters, int8_t *A, int8_t *B, int32_t *C, size_t M, size_t N, size_t K) {
  for(int i = 0; i < niters; ++i) {
    for(int mi=0; mi<(M/TILE_M); mi+=2) {
      for(int ni=0; ni<(N/TILE_N); ni+=2) {
        for(int ki=0; ki<(K/TILE_K); ++ki) {
          _tile_loadd(TMM0, A+mi*TILE_M*K+ki*TILE_K, K);
          _tile_loadd(TMM2, B+ki*(TILE_K/4)*(N*4)+ni*TILE_N*4, N*4);
          _tile_loadd(TMM1, A+(mi+1)*TILE_M*K+ki*TILE_K, K);
          _tile_loadd(TMM3, B+ki*(TILE_K/4)*(N*4)+(ni+1)*TILE_N*4, N*4);
        }
      }
    }
  }
  return 0;
}

// store bw utilization: 100%
int l1_bw_store(int niters, int8_t *A, int8_t *B, int32_t *C, size_t M, size_t N, size_t K) {
  for(int i = 0; i < niters; ++i) {
    for(int mi=0; mi<(M/TILE_M); mi+=2) {
      for(int ni=0; ni<(N/TILE_N); ni+=2) {
        _tile_zero(TMM0);
        _tile_zero(TMM1);
        _tile_zero(TMM2);
        _tile_zero(TMM3);

        _tile_stored(TMM0, (int8_t *)C+mi*TILE_M*N*sizeof(int32_t)+ni*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        _tile_stored(TMM1, (int8_t *)C+mi*TILE_M*N*sizeof(int32_t)+(ni+1)*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        _tile_stored(TMM2, (int8_t *)C+(mi+1)*TILE_M*N*sizeof(int32_t)+ni*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
        _tile_stored(TMM3, (int8_t *)C+(mi+1)*TILE_M*N*sizeof(int32_t)+(ni+1)*TILE_N*sizeof(int32_t), N*sizeof(int32_t));
      }
    }
  }
  return 0;
}

void l1_stats(const char* loop_order, const size_t M, const size_t N, const size_t K, 
              const double ops, const double elapsed) {
    double load_gbytes = (M/(TILE_M*2)) * (N/(TILE_N*2)) * (K/TILE_K) * (TILE_M*2*TILE_K+TILE_N*2*TILE_K) / 1e9;
    double store_gbytes = (M/(TILE_M*2)) * (N/(TILE_N*2)) * (4*TILE_M*TILE_K) / 1e9;
    double load_bw = load_gbytes / elapsed;
    double store_bw = store_gbytes / elapsed;

    printf("Read GB:     [%.2f] GB\n", load_gbytes);
    printf("Store GB:    [%.2f] GB\n", store_gbytes);
    printf("Read BW:     [%.2f] GB/s\n", load_bw);
    printf("Store BW:    [%.2f] GB/s\n", store_bw);

    printf("Read BW Utilization: [%.2f%%]\n", (load_bw / (L1_LOAD_BW*FREQ)) * 100.0); 
    printf("Store BW Utilization: [%.2f%%]\n", (store_bw / (L1_STORE_BW*FREQ)) * 100.0);   
}